module Regions

using LinearAlgebra
using ..FLEXINVERT
using NCDatasets
using ..FLEXINVERT.Settings: Files, Config
using ..FLEXINVERT.CoreTypes: Domain
using ..FLEXINVERT.Tools: grid_area, argsort_indices

"""
    get_regions!(dom::Domain, files::Files, cfg::Config)

Port of mod_regions:get_regions. Computes mapping nbox_xy, areas, and local time offsets.
Currently supports the non-aggregated case (cfg.regions=false). Region-file based
aggregation is planned via NCDatasets.
"""
function get_regions!(dom::Domain, files::Files, cfg::Config)
    if cfg.regions
        read_regions!(dom, files, cfg)
    else
        grid_boxes!(dom, cfg)
    end

    # Lat/lon vectors for region grid
    dom.reg_lon = [dom.rllx + (ix-1)*dom.rdx for ix in 1:dom.nxregrid]
    dom.reg_lat = [dom.rlly + (jy-1)*dom.rdy for jy in 1:dom.nyregrid]

    # Areas per box (land) and per ocean region
    dom.area_box = zeros(dom.nbox)
    if dom.nocn > 0
        dom.area_ocn = zeros(dom.nocn)
    else
        dom.area_ocn = Float64[]
    end

    for (j, lat) in enumerate(dom.reg_lat)
        area = grid_area(lat + 0.5*dom.rdy, dom.rdy, dom.rdx)
        for i in 1:dom.nxregrid
            idx = dom.nbox_xy[i, j]
            if idx > 0
                dom.area_box[idx] += area
            elseif dom.nocn > 0
                dom.area_ocn[-idx] += area
            end
        end
    end

    # Local time offset per region = median of cell offsets (hours)
    hloc_xy = [floor(Int, lon*24/360) for lon in dom.reg_lon]
    dom.hloc = zeros(dom.nbox)
    work = Float64[]
    for n in 1:dom.nbox
        empty!(work)
        for j in 1:dom.nyregrid, i in 1:dom.nxregrid
            dom.nbox_xy[i, j] == n || continue
            push!(work, float(hloc_xy[i]))
        end
        if !isempty(work)
            # median
            p = argsort_indices(work)
            dom.hloc[n] = work[p[cld(length(work), 2)]]
        end
    end
    return dom
end

"""
    grid_boxes!(dom::Domain, cfg::Config)

Non-aggregated grid: derive nbox_xy and lsm_box from land-sea mask (not yet available).
We set lsm_box=1 everywhere and respect cfg.inc_ocean to separate ocean as negative ids.
"""
function grid_boxes!(dom::Domain, cfg::Config)
    dom.nbox_xy = zeros(Int, dom.nxregrid, dom.nyregrid)
    dom.lsm_box = zeros(Int, dom.nxregrid, dom.nyregrid)
    n_water = 0
    n_invert = 0
    for j in 1:dom.nyregrid, i in 1:dom.nxregrid
        # use actual land-sea mask if available (threshold 0.01 as in Fortran)
        is_ocean = (!isempty(dom.lsm)) ? (dom.lsm[i, j] < 0.01) : false
        if (!cfg.inc_ocean) && is_ocean
            n_water -= 1
            dom.nbox_xy[i, j] = n_water
            dom.lsm_box[i, j] = 0
        else
            n_invert += 1
            dom.nbox_xy[i, j] = n_invert
            dom.lsm_box[i, j] = 1
        end
    end

    # Handle case where all boxes are zero (shouldn't happen but provides safety)
    if n_invert == 0 && n_water == 0
        @warn "No valid grid boxes found, treating all as land"
        for j in 1:dom.nyregrid, i in 1:dom.nxregrid
            n_invert += 1
            dom.nbox_xy[i, j] = n_invert
            dom.lsm_box[i, j] = 1
        end
    end

    dom.nbox = maximum(dom.nbox_xy)
    dom.nocn = n_water != 0 ? abs(minimum(dom.nbox_xy)) : 0
    return dom
end

"""
    read_regions!(dom::Domain, files::Files, cfg::Config)

Reads a NetCDF region mask dataset and fills dom.nbox_xy and dom.lsm_box like Fortran's read_regions.
"""
function read_regions!(dom::Domain, files::Files, cfg::Config)
    filename = files.file_regions
    varname  = files.varname_regs
    lonname  = files.lonname_regs
    latname  = files.latname_regs
    isfile(filename) || error("read_regions!: cannot find $(filename)")
    ds = NCDataset(filename)
    try
        lon = vec(ds[lonname][:])
        lat = vec(ds[latname][:])
        reg = Array(ds[varname][:])  # (nlon, nlat)
    finally
        close(ds)
    end
    # Find indices covering the inversion domain
    xres = lon[2] - lon[1]
    yres = lat[2] - lat[1]
    # grid centers expected at lon,lat; Fortran uses edges +/- 0.5*res
    ix1 = argmin(abs.(lon .- (dom.rllx + 0.5*dom.rdx)))
    ix2 = argmin(abs.(lon .- (dom.rurx - 0.5*dom.rdx)))
    jy1 = argmin(abs.(lat .- (dom.rlly + 0.5*dom.rdy)))
    jy2 = argmin(abs.(lat .- (dom.rury - 0.5*dom.rdy)))
    @assert ix2 - ix1 + 1 == dom.nxregrid "incompatible longitudinal extent"
    @assert jy2 - jy1 + 1 == dom.nyregrid "incompatible latitudinal extent"
    sub = reg[ix1:ix2, jy1:jy2]
    dom.nbox_xy = zeros(Int, size(sub))
    dom.lsm_box = zeros(Int, size(sub))
    n_water = 0
    n_invert = 0
    for val in unique(vec(sub))
        if any(sub .== val)
            if (val < 0) && (!cfg.inc_ocean)
                n_water -= 1
                dom.nbox_xy[sub .== val] .= n_water
            else
                n_invert += 1
                dom.nbox_xy[sub .== val] .= n_invert
            end
            if val < 0
                dom.lsm_box[sub .== val] .= 0
            else
                dom.lsm_box[sub .== val] .= 1
            end
        end
    end
    dom.nbox = maximum(dom.nbox_xy)
    dom.nocn = n_water != 0 ? abs(minimum(dom.nbox_xy)) : 0
    return dom
end


end # module
